from lexer import TokenType, Token

class ASTNode:
    pass

class Boolean(ASTNode):
    def __init__(self, value):
        self.value = value
        
    def __repr__(self):
        return f"Boolean({self.value})"

class UsingStatement(ASTNode):
    def __init__(self, module_name):
        self.module_name = module_name
        
    def __repr__(self):
        return f"UsingStatement({self.module_name})"

class ClassDef(ASTNode):
    def __init__(self, name, body):
        self.name = name
        self.body = body
        
    def __repr__(self):
        return f"ClassDef({self.name}, {self.body})"

class MethodDef(ASTNode):
    def __init__(self, name, params, body):
        self.name = name
        self.params = params
        self.body = body
        
    def __repr__(self):
        return f"MethodDef({self.name}, {self.params}, {self.body})"

class PrintStatement(ASTNode):
    def __init__(self, expr):
        self.expr = expr
        
    def __repr__(self):
        return f"PrintStatement({self.expr})"

class Variable(ASTNode):
    def __init__(self, name):
        self.name = name
        
    def __repr__(self):
        return f"Variable({self.name})"

class Number(ASTNode):
    def __init__(self, value):
        self.value = value
        
    def __repr__(self):
        return f"Number({self.value})"

class String(ASTNode):
    def __init__(self, value):
        self.value = value
        
    def __repr__(self):
        return f"String({self.value})"

class MethodCall(ASTNode):
    def __init__(self, obj_name, method_name, args):
        self.obj_name = obj_name
        self.method_name = method_name
        self.args = args
        
    def __repr__(self):
        return f"MethodCall({self.obj_name}.{self.method_name}, {self.args})"

class ObjectCreation(ASTNode):
    def __init__(self, class_name):
        self.class_name = class_name
        
    def __repr__(self):
        return f"ObjectCreation({self.class_name})"

class Parser:
    def __init__(self, lexer):
        self.lexer = lexer
        self.current_token = self.lexer.get_next_token()
        
    def error(self):
        raise Exception(f"Invalid syntax at position {self.lexer.pos}. Unexpected token: {self.current_token}")
        
    def eat(self, token_type):
        if self.current_token.type == token_type:
            self.current_token = self.lexer.get_next_token()
        else:
            self.error()
            
    def parse(self):
        statements = []
        while self.current_token.type != TokenType.EOF:
      
            if self.current_token.type == TokenType.USING:
                statements.append(self.parse_using())
            elif self.current_token.type == TokenType.CLASS:
                statements.append(self.parse_class())
            elif self.current_token.type == TokenType.DEF:
                statements.append(self.parse_method())
            elif self.current_token.type == TokenType.PRINT:
                statements.append(self.parse_print())
            elif self.current_token.type == TokenType.PRINTIN:
                statements.append(self.parse_printin())
            elif self.current_token.type == TokenType.INPUTON:
                statements.append(self.parse_inputon())
            elif self.current_token.type == TokenType.READFILE:
                statements.append(self.parse_readfile())
            elif self.current_token.type == TokenType.WRITEFILE:
                statements.append(self.parse_writefile())
            elif self.current_token.type == TokenType.IF:
                statements.append(self.parse_if())
            elif self.current_token.type == TokenType.IDENTIFIER:
                # Could be variable assignment or method call
                statements.append(self.parse_statement())
            elif self.current_token.type == TokenType.RBRACE:
                # End of block, let the caller handle it
                break
            else:
                self.error()
        return statements
        
    def parse_using(self):
        self.eat(TokenType.USING)
        module_name = self.current_token.value
        self.eat(TokenType.IDENTIFIER)
        self.eat(TokenType.SEMICOLON)
        return UsingStatement(module_name)
        
    def parse_class(self):
        self.eat(TokenType.CLASS)
        name = self.current_token.value
        self.eat(TokenType.IDENTIFIER)
        self.eat(TokenType.LBRACE)
        
        body = []
        while self.current_token.type != TokenType.RBRACE:
            if self.current_token.type == TokenType.DEF:
                body.append(self.parse_method())
            else:
                body.append(self.parse_statement())
                
        self.eat(TokenType.RBRACE)
        return ClassDef(name, body)
        
    def parse_method(self):
        self.eat(TokenType.DEF)
        name = self.current_token.value
        self.eat(TokenType.IDENTIFIER)
        self.eat(TokenType.LPAREN)
        
        params = []
        if self.current_token.type == TokenType.IDENTIFIER:
            params.append(self.current_token.value)
            self.eat(TokenType.IDENTIFIER)
            
            while self.current_token.type == TokenType.COMMA:
                self.eat(TokenType.COMMA)
                params.append(self.current_token.value)
                self.eat(TokenType.IDENTIFIER)
                
        self.eat(TokenType.RPAREN)
        self.eat(TokenType.LBRACE)
        
        body = []
        while self.current_token.type != TokenType.RBRACE:
            
            statements = self.parse()  # parse() returns a list of statements

            body.extend(statements)  # Use extend to flatten the list
                
        self.eat(TokenType.RBRACE)
       
        return MethodDef(name, params, body)
        
    def parse_print(self):
        self.eat(TokenType.PRINT)
        expr = self.parse_expression()
        self.eat(TokenType.SEMICOLON)
        return PrintStatement(expr)
        
    def parse_printin(self):
        self.eat(TokenType.PRINTIN)
        expr = self.parse_expression()
        self.eat(TokenType.SEMICOLON)
        return PrintStatement(expr)
        
    def parse_inputon(self):
        self.eat(TokenType.INPUTON)
        self.eat(TokenType.LPAREN)
        prompt = self.parse_expression()
        self.eat(TokenType.RPAREN)
        self.eat(TokenType.SEMICOLON)
        return ("inputon", prompt)
        
    def parse_readfile(self):
        self.eat(TokenType.READFILE)
        self.eat(TokenType.LPAREN)
        filename = self.parse_expression()
        self.eat(TokenType.RPAREN)
        self.eat(TokenType.SEMICOLON)
        return ("readfile", filename)
        
    def parse_writefile(self):
        self.eat(TokenType.WRITEFILE)
        self.eat(TokenType.LPAREN)
        filename = self.parse_expression()
        self.eat(TokenType.COMMA)
        content = self.parse_expression()
        self.eat(TokenType.RPAREN)
        self.eat(TokenType.SEMICOLON)
        return ("writefile", filename, content)
        
    def parse_statement(self):
        if self.current_token.type == TokenType.IDENTIFIER:
            name = self.current_token.value
            self.eat(TokenType.IDENTIFIER)
            
            if self.current_token.type == TokenType.EQUAL:
                self.eat(TokenType.EQUAL)
                value = self.parse_expression()
                self.eat(TokenType.SEMICOLON)
                return ("assignment", name, value)
            elif self.current_token.type == TokenType.DOT:
                # Object method call
                self.eat(TokenType.DOT)
                method_name = self.current_token.value
                self.eat(TokenType.IDENTIFIER)
                self.eat(TokenType.LPAREN)
                args = []
                if self.current_token.type != TokenType.RPAREN:
                    args.append(self.parse_expression())
                    while self.current_token.type == TokenType.COMMA:
                        self.eat(TokenType.COMMA)
                        args.append(self.parse_expression())
                self.eat(TokenType.RPAREN)
                self.eat(TokenType.SEMICOLON)
                
                return MethodCall(name, method_name, args)
                
            elif self.current_token.type == TokenType.LPAREN:
                # Global function call
                self.eat(TokenType.LPAREN)
                args = []
                if self.current_token.type != TokenType.RPAREN:
                    args.append(self.parse_expression())
                    while self.current_token.type == TokenType.COMMA:
                        self.eat(TokenType.COMMA)
                        args.append(self.parse_expression())
                self.eat(TokenType.RPAREN)
                self.eat(TokenType.SEMICOLON)
                
                return MethodCall(None, name, args)
            else:
                self.error()
        else:
            self.error()
            
    def parse_expression(self):
        # First parse a simple expression
        node = self.parse_simple_expression()
        
        # Then check for binary operations
        while self.current_token.type in (TokenType.PLUS, TokenType.MINUS):
            op = self.current_token.type
            self.eat(op)
            right = self.parse_simple_expression()
            node = ("+" if op == TokenType.PLUS else "-", node, right)
            
        return node
        
    def parse_if(self):
        self.eat(TokenType.IF)
        
        # Parse condition expression with comparison operators
        left = self.parse_expression()
        
        # Check for comparison operators
        if self.current_token.type in (TokenType.EQ, TokenType.NE, TokenType.LT, 
                                 TokenType.GT, TokenType.LTE, TokenType.GTE):
            op = self.current_token.type
            op_str = {
                TokenType.EQ: "==",
                TokenType.NE: "!=",
                TokenType.LT: "<",
                TokenType.GT: ">",
                TokenType.LTE: "<=",
                TokenType.GTE: ">="
            }[op]
            self.eat(op)
            right = self.parse_expression()
            condition = (op_str, left, right)
        else:
            condition = left
            
        self.eat(TokenType.LBRACE)
        
        body = []
        while self.current_token.type != TokenType.RBRACE:
            statements = self.parse()  # parse() returns a list of statements
            body.extend(statements)  # Use extend to flatten the list
            
        self.eat(TokenType.RBRACE)
        
        # Optional else clause
        else_body = []
        if self.current_token.type == TokenType.ELSE:
            self.eat(TokenType.ELSE)
            self.eat(TokenType.LBRACE)
            
            while self.current_token.type != TokenType.RBRACE:
                statements = self.parse()  # parse() returns a list of statements
                else_body.extend(statements)  # Use extend to flatten the list
                
            self.eat(TokenType.RBRACE)
            
        return ("if", condition, body, else_body)
        
    def parse_simple_expression(self):
        if self.current_token.type == TokenType.NUMBER:
            node = Number(self.current_token.value)
            self.eat(TokenType.NUMBER)
            return node
        elif self.current_token.type == TokenType.STRING:
            node = String(self.current_token.value)
            self.eat(TokenType.STRING)
            return node
        elif self.current_token.type == TokenType.TRUE:
            self.eat(TokenType.TRUE)
            return Boolean(True)
        elif self.current_token.type == TokenType.FALSE:
            self.eat(TokenType.FALSE)
            return Boolean(False)
        elif self.current_token.type == TokenType.IDENTIFIER:
            name = self.current_token.value
            self.eat(TokenType.IDENTIFIER)
            
            if self.current_token.type == TokenType.LPAREN:
                # Method call
                self.eat(TokenType.LPAREN)
                args = []
                if self.current_token.type != TokenType.RPAREN:
                    args.append(self.parse_expression())
                    while self.current_token.type == TokenType.COMMA:
                        self.eat(TokenType.COMMA)
                        args.append(self.parse_expression())
                self.eat(TokenType.RPAREN)
                return MethodCall(None, name, args)
            elif name == "new":
                # Object creation
                class_name = self.current_token.value
                self.eat(TokenType.IDENTIFIER)
                self.eat(TokenType.LPAREN)
                self.eat(TokenType.RPAREN)
                return ObjectCreation(class_name)
            else:
                # Simple variable
                return Variable(name)
        else:
            self.error()
